# -*- coding: utf-8 -*-
# PyTorch_CIFAR10_GPU.py
"""
PyTorch GPU加速实验：CIFAR-10图像分类
使用PyTorch构建CNN模型，利用GPU加速训练，对CIFAR-10数据集进行分类
针对容器环境修改了文件保存路径
"""

import torch
import torch.nn as nn
import torch.nn.functional as F
import torch.optim as optim
import torchvision
import torchvision.transforms as transforms
import matplotlib.pyplot as plt
import numpy as np
import time
from torch.utils.tensorboard import SummaryWriter
import datetime
import os

# 设置工作目录为挂载的卷
WORKSPACE_DIR = "/home/vscode/workspace"
os.chdir(WORKSPACE_DIR)

# 检查GPU可用性
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
print(f"使用的设备: {device}")

if torch.cuda.is_available():
    print(f"GPU名称: {torch.cuda.get_device_name(0)}")
    print(f"GPU内存: {torch.cuda.get_device_properties(0).total_memory / 1024**3:.2f} GB")

# 定义卷积神经网络模型
class CNN(nn.Module):
    def __init__(self):
        super(CNN, self).__init__()
        self.conv1 = nn.Conv2d(3, 32, 3, padding=1)
        self.conv2 = nn.Conv2d(32, 64, 3, padding=1)
        self.conv3 = nn.Conv2d(64, 128, 3, padding=1)
        self.pool = nn.MaxPool2d(2, 2)
        self.fc1 = nn.Linear(128 * 4 * 4, 512)
        self.fc2 = nn.Linear(512, 256)
        self.fc3 = nn.Linear(256, 10)
        self.dropout = nn.Dropout(0.25)
        self.batch_norm1 = nn.BatchNorm2d(32)
        self.batch_norm2 = nn.BatchNorm2d(64)
        self.batch_norm3 = nn.BatchNorm2d(128)

    def forward(self, x):
        x = self.pool(F.relu(self.batch_norm1(self.conv1(x))))
        x = self.pool(F.relu(self.batch_norm2(self.conv2(x))))
        x = self.pool(F.relu(self.batch_norm3(self.conv3(x))))
        x = x.view(-1, 128 * 4 * 4)
        x = self.dropout(F.relu(self.fc1(x)))
        x = self.dropout(F.relu(self.fc2(x)))
        x = self.fc3(x)
        return x

# 数据预处理和加载
transform = transforms.Compose([
    transforms.RandomHorizontalFlip(),
    transforms.RandomCrop(32, padding=4),
    transforms.ToTensor(),
    transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))
])

# 创建数据目录
data_dir = os.path.join(WORKSPACE_DIR, 'data')
os.makedirs(data_dir, exist_ok=True)

# 加载CIFAR-10数据集
trainset = torchvision.datasets.CIFAR10(root=data_dir, train=True,
                                        download=True, transform=transform)
trainloader = torch.utils.data.DataLoader(trainset, batch_size=128,
                                          shuffle=True, num_workers=2)

testset = torchvision.datasets.CIFAR10(root=data_dir, train=False,
                                       download=True, transform=transform)
testloader = torch.utils.data.DataLoader(testset, batch_size=128,
                                         shuffle=False, num_workers=2)

# 类别名称 - 使用英文避免字体问题
classes = ('airplane', 'automobile', 'bird', 'cat', 'deer', 
           'dog', 'frog', 'horse', 'ship', 'truck')

# 初始化模型、损失函数和优化器
model = CNN().to(device)
criterion = nn.CrossEntropyLoss()
optimizer = optim.Adam(model.parameters(), lr=0.001, weight_decay=1e-5)
scheduler = optim.lr_scheduler.StepLR(optimizer, step_size=10, gamma=0.5)

# 创建日志目录
log_dir = os.path.join(WORKSPACE_DIR, "runs/cifar10_experiment_" + datetime.datetime.now().strftime("%Y%m%d-%H%M%S"))
os.makedirs(os.path.dirname(log_dir), exist_ok=True)

# 设置TensorBoard记录
writer = SummaryWriter(log_dir)

# 修改后的训练函数，增加device参数
def train_model(model, trainloader, criterion, optimizer, num_epochs=25, device=None, verbose=True):
    if device is None:
        device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
    
    since = time.time()
    best_acc = 0.0
    train_losses = []
    train_accs = []
    
    for epoch in range(num_epochs):
        if verbose:
            print(f'第 {epoch+1}/{num_epochs} 轮训练')
            print('-' * 10)
        
        # 训练阶段
        model.train()
        running_loss = 0.0
        running_corrects = 0
        
        for i, data in enumerate(trainloader, 0):
            inputs, labels = data
            inputs, labels = inputs.to(device), labels.to(device)
            
            optimizer.zero_grad()
            
            outputs = model(inputs)
            _, preds = torch.max(outputs, 1)
            loss = criterion(outputs, labels)
            
            loss.backward()
            optimizer.step()
            
            running_loss += loss.item() * inputs.size(0)
            running_corrects += torch.sum(preds == labels.data)
            
            if verbose and i % 100 == 99:  # 每100个batch打印一次
                print(f'[{epoch + 1}, {i + 1:5d}] 损失: {running_loss / (100 * inputs.size(0)):.3f}')
                
        epoch_loss = running_loss / len(trainset)
        epoch_acc = running_corrects.double() / len(trainset)
        
        train_losses.append(epoch_loss)
        train_accs.append(epoch_acc.cpu().numpy())
        
        if verbose:
            print(f'训练损失: {epoch_loss:.4f} 准确率: {epoch_acc:.4f}')
        
        # 记录到TensorBoard
        if verbose:
            writer.add_scalar(f'{device}_训练损失', epoch_loss, epoch)
            writer.add_scalar(f'{device}_训练准确率', epoch_acc, epoch)
        
        # 更新学习率
        scheduler.step()
        
        # 每个epoch结束后在测试集上验证
        test_loss, test_acc = evaluate_model(model, testloader, criterion, device)
        if verbose:
            writer.add_scalar(f'{device}_测试损失', test_loss, epoch)
            writer.add_scalar(f'{device}_测试准确率', test_acc, epoch)
            print(f'测试损失: {test_loss:.4f} 准确率: {test_acc:.4f}')
        
        # 保存最佳模型
        if test_acc > best_acc:
            best_acc = test_acc
            model_path = os.path.join(WORKSPACE_DIR, f'best_model_{device}.pth')
            torch.save(model.state_dict(), model_path)
    
    time_elapsed = time.time() - since
    if verbose:
        print(f'训练完成于 {time_elapsed // 60:.0f}分 {time_elapsed % 60:.0f}秒')
        print(f'最佳测试准确率: {best_acc:.4f}')
    
    return model, train_losses, train_accs

# 修改后的评估函数，增加device参数
def evaluate_model(model, testloader, criterion, device=None):
    if device is None:
        device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
        
    model.eval()
    running_loss = 0.0
    running_corrects = 0
    
    with torch.no_grad():
        for data in testloader:
            images, labels = data
            images, labels = images.to(device), labels.to(device)
            
            outputs = model(images)
            _, preds = torch.max(outputs, 1)
            loss = criterion(outputs, labels)
            
            running_loss += loss.item() * images.size(0)
            running_corrects += torch.sum(preds == labels.data)
    
    test_loss = running_loss / len(testset)
    test_acc = running_corrects.double() / len(testset)
    
    return test_loss, test_acc

# 可视化一些预测结果
def visualize_predictions(model, testloader, num_images=12):
    model.eval()
    images_so_far = 0
    fig = plt.figure(figsize=(12, 8))
    
    with torch.no_grad():
        for data in testloader:
            images, labels = data
            images, labels = images.to(device), labels.to(device)
            
            outputs = model(images)
            _, preds = torch.max(outputs, 1)
            
            for j in range(images.size()[0]):
                images_so_far += 1
                ax = plt.subplot(num_images//4, 4, images_so_far)
                ax.axis('off')
                
                # 反标准化图像
                img = images.cpu().data[j] / 2 + 0.5
                img = img.numpy().transpose((1, 2, 0))
                
                ax.imshow(img)
                ax.set_title(f'Pred: {classes[preds[j]]}\nTrue: {classes[labels[j]]}',
                            color=('green' if preds[j] == labels[j] else 'red'))
                
                if images_so_far == num_images:
                    plt.tight_layout()
                    save_path = os.path.join(WORKSPACE_DIR, 'predictions.png')
                    plt.savefig(save_path, dpi=300, bbox_inches='tight')
                    plt.show()
                    return

# 绘制训练曲线
def plot_training_curves(train_losses, train_accs):
    fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(12, 4))
    
    # 绘制损失曲线
    ax1.plot(train_losses)
    ax1.set_title('Training Loss')
    ax1.set_xlabel('Epoch')
    ax1.set_ylabel('Loss')
    
    # 绘制准确率曲线
    ax2.plot(train_accs)
    ax2.set_title('Training Accuracy')
    ax2.set_xlabel('Epoch')
    ax2.set_ylabel('Accuracy')
    
    plt.tight_layout()
    save_path = os.path.join(WORKSPACE_DIR, 'training_curves.png')
    plt.savefig(save_path, dpi=300, bbox_inches='tight')
    plt.show()

# 性能对比函数
def compare_cpu_gpu_performance(num_epochs=3):
    """比较CPU和GPU的训练性能"""
    results = []
    
    for device_name in ['cpu', 'cuda']:
        if device_name == 'cuda' and not torch.cuda.is_available():
            print("CUDA不可用，跳过GPU测试")
            continue
            
        print(f"\n=== 使用{device_name.upper()}进行训练 ===")
        current_device = torch.device(device_name)
        
        # 重新初始化模型和优化器
        model = CNN().to(current_device)
        criterion = nn.CrossEntropyLoss()
        optimizer = optim.Adam(model.parameters(), lr=0.001, weight_decay=1e-5)
        scheduler = optim.lr_scheduler.StepLR(optimizer, step_size=10, gamma=0.5)
        
        # 记录开始时间
        start_time = time.time()
        
        # 训练模型
        model, train_losses, train_accs = train_model(
            model, trainloader, criterion, optimizer, num_epochs, current_device, verbose=False
        )
        
        # 记录结束时间
        end_time = time.time()
        total_time = end_time - start_time
        
        # 评估模型
        test_loss, test_acc = evaluate_model(model, testloader, criterion, current_device)
        
        results.append({
            'device': device_name,
            'total_time': total_time,
            'avg_epoch_time': total_time / num_epochs,
            'test_accuracy': test_acc
        })
        
        print(f"{device_name.upper()}训练完成 - 总时间: {total_time:.2f}秒, "
              f"平均每轮: {total_time/num_epochs:.2f}秒, 测试准确率: {test_acc:.4f}")
    
    return results

# 性能对比可视化
def plot_performance_comparison(results):
    """可视化性能对比结果"""
    if not results:
        print("没有性能数据可可视化")
        return
        
    devices = [r['device'] for r in results]
    times = [r['avg_epoch_time'] for r in results]
    accuracies = [r['test_accuracy'].cpu().numpy() for r in results]  # 转换为numpy数组
    
    fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(12, 5))
    
    # 绘制时间对比
    bars = ax1.bar(devices, times, color=['blue', 'green'])
    ax1.set_title('平均每轮训练时间对比')
    ax1.set_ylabel('时间 (秒)')
    ax1.set_xlabel('设备')
    
    # 在柱状图上添加数值标签
    for bar, time_val in zip(bars, times):
        height = bar.get_height()
        ax1.text(bar.get_x() + bar.get_width()/2., height + 0.01,
                f'{time_val:.2f}s', ha='center', va='bottom')
    
    # 绘制准确率对比
    bars = ax2.bar(devices, accuracies, color=['blue', 'green'])
    ax2.set_title('测试准确率对比')
    ax2.set_ylabel('准确率')
    ax2.set_xlabel('设备')
    ax2.set_ylim(0, 1)
    
    # 在柱状图上添加数值标签
    for bar, acc in zip(bars, accuracies):
        height = bar.get_height()
        ax2.text(bar.get_x() + bar.get_width()/2., height + 0.01,
                f'{acc:.4f}', ha='center', va='bottom')
    
    plt.tight_layout()
    save_path = os.path.join(WORKSPACE_DIR, 'performance_comparison.png')
    plt.savefig(save_path, dpi=300, bbox_inches='tight')
    plt.show()
    
    # 计算加速比
    if len(results) == 2:
        cpu_time = results[0]['avg_epoch_time']
        gpu_time = results[1]['avg_epoch_time']
        speedup = cpu_time / gpu_time
        print(f"\nGPU相对于CPU的加速比: {speedup:.2f}x")

# 主函数
def main():
    global model, trainloader, testloader, criterion, optimizer
    
    # 选择是否要运行性能对比
    run_comparison = input("是否要运行CPU/GPU性能对比? (y/n): ").lower().startswith('y')
    
    if run_comparison:
        print("开始CPU/GPU性能对比...")
        results = compare_cpu_gpu_performance(num_epochs=3)
        plot_performance_comparison(results)
    
    print("开始训练CNN模型...")
    
    # 训练模型
    model, train_losses, train_accs = train_model(
        model, trainloader, criterion, optimizer, num_epochs=30, device=device
    )
    
    # 绘制训练曲线
    plot_training_curves(train_losses, train_accs)
    
    # 加载最佳模型
    model_path = os.path.join(WORKSPACE_DIR, f'best_model_{device}.pth')
    # 使用weights_only=True提高安全性（如果PyTorch版本支持）
    try:
        model.load_state_dict(torch.load(model_path, weights_only=True))
    except:
        model.load_state_dict(torch.load(model_path))
    
    # 最终评估
    test_loss, test_acc = evaluate_model(model, testloader, criterion, device)
    print(f'最终测试结果 - 损失: {test_loss:.4f}, 准确率: {test_acc:.4f}')
    
    # 可视化一些预测结果
    try:
        visualize_predictions(model, testloader)
    except KeyboardInterrupt:
        print("可视化被用户中断")
    except Exception as e:
        print(f"可视化过程中出现错误: {e}")
    
    # 关闭TensorBoard writer
    writer.close()
    
    print("实验完成! 可以使用 'tensorboard --logdir=runs' 查看训练详情")

if __name__ == '__main__':
    main()